package com.example.handpose_front.utils.OKHttpUtil;

import android.util.Log;

import androidx.annotation.NonNull;
import androidx.camera.core.ImageProxy;

import com.google.gson.Gson;
import com.google.gson.annotations.SerializedName;

import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.Response;
import okhttp3.WebSocket;
import okhttp3.WebSocketListener;
import okio.ByteString;

public class GestureWebSocketClient {
    private final OkHttpClient client;
    private WebSocket webSocket;
    private final String wsUrl;
    private final Gson gson = new Gson();
    private boolean isConnected = false;
    private final ScheduledExecutorService reconnectExecutor = Executors.newSingleThreadScheduledExecutor();
    private GestureCallback callback;

    // 心跳检测间隔(秒)
    private static final int HEARTBEAT_INTERVAL = 10;
    private static final int RECONNECT_DELAY = 5; // 重连等待时间(秒)

    public interface GestureCallback {
        void onResult(GestureResult result);
        void onError(String error);
    }

    public static class GestureResult {
        @SerializedName("gesture")
        public String gesture;
        @SerializedName("confidence")
        public float confidence;
        @SerializedName("landmarks")
        public float[][] landmarks; // 手部关键点坐标

        @Override
        public String toString() {
            return "Gesture: " + gesture + " (Confidence: " + confidence + ")";
        }
    }

    public GestureWebSocketClient(String serverUrl) {
        this.wsUrl = serverUrl;
        this.client = new OkHttpClient.Builder()
                .pingInterval(HEARTBEAT_INTERVAL, TimeUnit.SECONDS)
                .retryOnConnectionFailure(true)
                .build();
    }

    public void connect(@NonNull GestureCallback callback) {
        this.callback = callback;
        Request request = new Request.Builder()
                .url(wsUrl)
                .build();

        webSocket = client.newWebSocket(request, new WebSocketListener() {
            @Override
            public void onOpen(WebSocket webSocket, Response response) {
                isConnected = true;
                Log.i("WebSocket", "Connected to server");
            }

            @Override
            public void onMessage(WebSocket webSocket, String text) {
                // 文本协议解析
                try {
                    GestureResult result = gson.fromJson(text, GestureResult.class);
                    callback.onResult(result);
                } catch (Exception e) {
                    callback.onError("Failed to parse text message: " + e.getMessage());
                }
            }

            @Override
            public void onMessage(WebSocket webSocket, ByteString bytes) {
                // 二进制协议解析
                String responseString = parseBinaryResponse(bytes);
                if (responseString != null) {
                    try {
                        GestureResult result = gson.fromJson(responseString, GestureResult.class);
                        callback.onResult(result);
                    } catch (Exception e) {
                        callback.onError("Failed to parse binary message: " + e.getMessage());
                    }
                }
            }

            @Override
            public void onClosed(WebSocket webSocket, int code, String reason) {
                isConnected = false;
                Log.w("WebSocket", "Connection closed: " + reason);
                scheduleReconnect();
            }

            @Override
            public void onFailure(WebSocket webSocket, Throwable t, Response response) {
                isConnected = false;
                Log.e("WebSocket", "Connection failed", t);
                callback.onError("Connection error: " + t.getMessage());
                scheduleReconnect();
            }
        });
    }

    private String parseBinaryResponse(ByteString bytes) {
        try {
            return bytes.utf8(); // 直接转成字符串
        } catch (Exception e) {
            Log.e("WebSocket", "Failed to parse message as string", e);
            return null;
        }
    }

    private void scheduleReconnect() {
        reconnectExecutor.schedule(() -> {
            if (!isConnected) {
                Log.i("WebSocket", "Attempting to reconnect...");
                connect(callback);
            }
        }, RECONNECT_DELAY, TimeUnit.SECONDS);
    }

    // 发送优化后的图像帧
    public void sendImageFrame(ImageProxy imageProxy) {
        if (!isConnected || webSocket == null) {
            Log.w("WebSocket", "Cannot send frame - not connected");
            return;
        }

        try {
            ByteBuffer yBuffer = imageProxy.getPlanes()[0].getBuffer();
            ByteBuffer uBuffer = imageProxy.getPlanes()[1].getBuffer();
            ByteBuffer vBuffer = imageProxy.getPlanes()[2].getBuffer();

            // YUV420转NV21（零拷贝优化）
            byte[] nv21 = YUV420toNV21(yBuffer, uBuffer, vBuffer,
                    imageProxy.getWidth(), imageProxy.getHeight());

            // 添加协议头
            byte[] packet = buildBinaryPacket(nv21, imageProxy.getWidth(), imageProxy.getHeight());

            webSocket.send(ByteString.of(packet));
        } catch (Exception e) {
            Log.e("WebSocket", "Error sending frame", e);
            callback.onError("Frame send error: " + e.getMessage());
        }
    }

    private byte[] YUV420toNV21(ByteBuffer yBuffer, ByteBuffer uBuffer, ByteBuffer vBuffer,
                                int width, int height) {
        // 实现YUV420到NV21的转换
        int ySize = width * height;
        int uvSize = ySize / 4;

        byte[] nv21 = new byte[ySize + uvSize * 2];

        // 复制Y数据
        yBuffer.get(nv21, 0, ySize);

        // 交错UV数据 (NV21格式是VU交替)
        for (int i = 0; i < uvSize; i++) {
            nv21[ySize + i * 2] = vBuffer.get();     // V
            nv21[ySize + i * 2 + 1] = uBuffer.get(); // U
        }

        return nv21;
    }

    private byte[] buildBinaryPacket(byte[] imageData, int width, int height) {
        ByteBuffer buffer = ByteBuffer.allocate(16 + imageData.length)
                .order(ByteOrder.BIG_ENDIAN)
                .putLong(System.currentTimeMillis()) // 8字节时间戳
                .putInt(width)                     // 4字节宽度
                .putInt(height)                    // 4字节高度
                .put(imageData);                   // 图像数据

        return buffer.array();
    }

    public void disconnect() {
        if (webSocket != null) {
            webSocket.close(1000, "Normal closure");
        }
        reconnectExecutor.shutdown();
    }

    public boolean isConnected() {
        return isConnected;
    }
}